
#include "EF_Solver1D_TDMA.h"

#include "../PhysicalField/PhysicalField.h"
#include "../Field/Field1D.h"
#include "../Diagnostic/Diagnostic1D.h"

EF_Solver1D_TDMA::EF_Solver1D_TDMA(PicParams* params_, ParallelMPI* pmpi_, Grid* grid_, PhysicalField* fields_, Diagnostic* diag_) :
    Solver1D(params_, pmpi_, grid_, fields_, diag_)
{
    ParallelMPI1D* pmpi1D = static_cast<ParallelMPI1D*>(pmpi);

    dx = params->cell_length[0];
    dx_inv_ = 1.0 / dx;
    dx_sq = dx * dx;
    nx = params->dims_space_global[0]+1;
    nx_source_left = params->nx_source_left;

    if(bc_em_type_x[0] == "Dirichlet")
    {
        bc_x_left = 1;
        bc_e_value_left = bc_em_value_x[0];
    }
    else if(bc_em_type_x[0] == "Neumann")
    {
        bc_x_left = 2;
        bc_e_derivative_left = bc_em_value_x[0];
    }
    else if(bc_em_type_x[0] == "Neumann_self_consistent")
    {
        bc_x_left = 2;
        bc_e_derivative_left = bc_em_value_x[0];
    }


    if(bc_em_type_x[1] == "Dirichlet")
    {
        bc_x_right = 1;
        bc_e_value_right = bc_em_value_x[1];
    }
    else if(bc_em_type_x[1] == "Neumann")
    {
        bc_x_right = 2;
        bc_e_derivative_right = bc_em_value_x[1];
    }
    else if(bc_em_type_x[1] == "Neumann_self_consistent")
    {
        bc_x_right = 2;
        bc_e_derivative_right = bc_em_value_x[1];   
    }


    if(pmpi1D->isMaster())
    {
        initTDMA();
    }

    
}

EF_Solver1D_TDMA::~EF_Solver1D_TDMA()
{
}

void EF_Solver1D_TDMA::operator()(ParallelMPI* pmpi, PhysicalField* fields, Diagnostic* diag)
{
    ParallelMPI1D* pmpi1D = static_cast<ParallelMPI1D*>(pmpi);

    Field1D<double>* Ex1D           = static_cast<Field1D<double>*>(fields->Ex_);
    Field1D<double>* rho1D          = static_cast<Field1D<double>*>(fields->rho_);
    Field1D<double>* phi1D          = static_cast<Field1D<double>*>(fields->phi_);
    Field1D<double>* rho1D_global   = static_cast<Field1D<double>*>(fields->rho_global);
    Field1D<double>* phi1D_global   = static_cast<Field1D<double>*>(fields->phi_global);
    Field1D<double>* Ex1D_global    = static_cast<Field1D<double>*>(fields->Ex_global);

    if(pmpi1D->isMaster())
    {
        solve_TDMA(rho1D_global, phi1D_global, diag);
        solve_Ex(phi1D_global, Ex1D_global);
    }

    pmpi1D->barrier();

    pmpi1D->scatterField(Ex1D_global, Ex1D);
    pmpi1D->scatterField(phi1D_global, phi1D);

}


void EF_Solver1D_TDMA::initTDMA()
{
    a = new double[nx - nx_source_left];
    b = new double[nx - nx_source_left];
    c = new double[nx - nx_source_left];
    f = new double[nx - nx_source_left];
    e = new double[nx - nx_source_left];
    d = new double[nx - nx_source_left];
    for(int i =1; i < nx-1-nx_source_left; i++)
    {
        a[i] = 1.0;
        b[i] = -2.0;
        c[i] = 1.0;
    }

    if(bc_x_left == 1){
        a[0] = 0.0;
        b[0] = 1.0;
        c[0] = 0.0;
    }
    else if(bc_x_left == 2){
        a[0] = 0.0;
        b[0] = 1.0;
        c[0] = -1.0;
    }

    if(bc_x_right == 1){
        a[nx-1-nx_source_left] = 0.0;
        b[nx-1-nx_source_left] = 1.0;
        c[nx-1-nx_source_left] = 0.0;
    }
    else if(bc_x_right == 2){
        a[nx-1-nx_source_left] = -1.0;
        b[nx-1-nx_source_left] = 1.0;
        c[nx-1-nx_source_left] = 0.0;
    }

}


void EF_Solver1D_TDMA::solve_TDMA(Field<double>* rho, Field<double>* phi, Diagnostic* diag)
{
    Field1D<double>* rho1D = static_cast<Field1D<double>*>(rho);
    Field1D<double>* phi1D = static_cast<Field1D<double>*>(phi);

    Diagnostic1D* diag1D = static_cast<Diagnostic1D*>(diag);

    if(bc_em_type_x[0] == "Neumann_self_consistent")
    {
        bc_e_derivative_left = 0.0;
        for(int ispec = 0; ispec < diag1D->n_species; ispec++)
        {
            bc_e_derivative_left += diag1D->sigma_left[ispec];
        }
        bc_e_derivative_left *= const_ephi0_inv;
    }

    if(bc_em_type_x[1] == "Neumann_self_consistent")
    {
        bc_e_derivative_right = 0.0;
        for(int ispec = 0; ispec < diag1D->n_species; ispec++)
        {
            bc_e_derivative_right += diag1D->sigma_right[ispec];
        }
        bc_e_derivative_right *= -const_ephi0_inv;
    }

    //The boundary value can be changed with time
    for(int i = 1; i < nx-1-nx_source_left; i++)
    {
        f[i] = -dx_sq * const_ephi0_inv * (*rho1D)(i+nx_source_left);
    }

    if(bc_x_left == 1){
        f[0] = bc_e_value_left;
    }
    else if(bc_x_left == 2){
        f[0] = bc_e_derivative_left * dx;
    }

    if(bc_x_right == 1){
        f[nx-1-nx_source_left] = bc_e_value_right;
    }
    else if(bc_x_right == 2){
        f[nx-1-nx_source_left] = -bc_e_derivative_right * dx;
    }

    e[0] = c[0] / b[0];
    d[0] = f[0] / b[0];
    for(int i =1; i < nx-1-nx_source_left; i++)
    {
        e[i] = c[i] / ( b[i] - a[i] * e[i-1] );
        d[i] = ( f[i] -a[i] * d[i-1] ) / ( b[i] - a[i] * e[i-1] );
    }

    (*phi1D)(nx-1) = ( f[nx-1-nx_source_left] - a[nx-1-nx_source_left] * d[nx-2-nx_source_left] ) / ( b[nx-1-nx_source_left] - a[nx-1-nx_source_left] * e[nx-2-nx_source_left] );
    for(int i = nx-2-nx_source_left; i >= 0; i--)
    {
        (*phi1D)(i+nx_source_left) = d[i] - e[i] * (*phi1D)(i+1+nx_source_left);
    }

    for(int i=0; i<nx_source_left; i++)
    {
        (*phi1D)(i) = (*phi1D)(nx_source_left);
        //cout<<(*phi1D)(i)<<endl;
    }


}


void EF_Solver1D_TDMA::solve_Ex(Field<double>* phi, Field<double>* Ex)
{
    Field1D<double>* phi1D = static_cast<Field1D<double>*>(phi);
    Field1D<double>* Ex1D = static_cast<Field1D<double>*>(Ex);


    for(int i = 1; i < nx-1; i++)
    {
        (*Ex1D)(i) = - ((*phi1D)(i+1) - (*phi1D)(i-1)) *0.5 * dx_inv_;
    }

    (*Ex1D)(0) = -(-3.0 * (*phi1D)(0) + 4.0 * (*phi1D)(1) - (*phi1D)(2)) *0.5 * dx_inv_;
    (*Ex1D)(nx-1) = -((*phi1D)(nx-3) - 4.0 * (*phi1D)(nx-2) + 3.0 * (*phi1D)(nx-1)) *0.5 * dx_inv_;

}


//no source region for electric field
void EF_Solver1D_TDMA::initTDMA_org()
{
    a = new double[nx];
    b = new double[nx];
    c = new double[nx];
    f = new double[nx];
    e = new double[nx];
    d = new double[nx];
    for(int i =1; i < nx-1; i++)
    {
        a[i] = 1.0;
        b[i] = -2.0;
        c[i] = 1.0;
    }
    if(bc_x_left == 1){
        a[0] = 0.0;
        b[0] = 1.0;
        c[0] = 0.0;
    }
    else if(bc_x_left == 2){
        a[0] = 0.0;
        b[0] = 1.0;
        c[0] = -1.0;
    }
    if(bc_x_right == 1){
        a[nx-1] = 0.0;
        b[nx-1] = 1.0;
        c[nx-1] = 0.0;
    }
    else if(bc_x_right == 2){
        a[nx-1] = -1.0;
        b[nx-1] = 1.0;
        c[nx-1] = 0.0;
    }
}

//no source region for electric field
void EF_Solver1D_TDMA::solve_TDMA_org(Field<double>* rho, Field<double>* phi)
{
    Field1D<double>* rho1D = static_cast<Field1D<double>*>(rho);
    Field1D<double>* phi1D = static_cast<Field1D<double>*>(phi);
    //The boundary value can be changed with time
    for(int i = 1; i < nx-1; i++)
    {
        f[i] = -dx_sq * const_ephi0_inv * (*rho1D)(i);
    }
    if(bc_x_left == 1){
        f[0] = bc_e_value_left;
    }
    else if(bc_x_left == 2){
        f[0] = bc_e_derivative_left;
    }
    if(bc_x_right == 1){
        f[nx-1] = bc_e_value_right;
    }
    else if(bc_x_right == 2){
        f[nx-1] = -bc_e_derivative_right;
    }
    e[0] = c[0] / b[0];
    d[0] = f[0] / b[0];
    for(int i =1; i < nx-1; i++)
    {
        e[i] = c[i] / ( b[i] - a[i] * e[i-1] );
        d[i] = ( f[i] -a[i] * d[i-1] ) / ( b[i] - a[i] * e[i-1] );
    }
    (*phi1D)(nx-1) = ( f[nx-1] - a[nx-1] * d[nx-2] ) / ( b[nx-1] - a[nx-1] * e[nx-2] );
    for(int i = nx-2; i >= 0; i--)
    {
        (*phi1D)(i) = d[i] - e[i] * (*phi1D)(i+1);
    }
}
